# Wigner 转动符号计算
# 参考： https://zhuanlan.zhihu.com/p/22974240583

import sympy as sp

g1,b1,g2,b2 = sp.symbols('g1,b1,g2,b2',real=True,positive=True)
L1 = sp.Matrix([[g1,g1*b1,0,0],[g1*b1,g1,0,0],[0,0,1,0],[0,0,0,1]])
L2 = sp.Matrix([[g2,0,g2*b2,0],[0,1,0,0],[g2*b2,0,g2,0],[0,0,0,1]])

L1_inv = sp.Matrix([[g1,-g1*b1,0,0],[-g1*b1,g1,0,0],[0,0,1,0],[0,0,0,1]])
L2_inv = sp.Matrix([[g2,0,-g2*b2,0],[0,1,0,0],[-g2*b2,0,g2,0],[0,0,0,1]])

test_data_dict = {b1:0.5,b2:0.2,g1:1.1547,g2:1.0206}

# 两次Lorentz变换
A = L2_inv @ L1_inv
print("组合变换矩阵A:")
print(A)
print(A.subs(test_data_dict).evalf())

c,s = sp.symbols('c,s',real=True)
Q = sp.Matrix([[1,0,0,0],[0,c,s,0],[0,-s,c,0],[0,0,0,1]])
Q_inv = Q.T

# A = P Q
P = A @ Q_inv
print("其中纯Lorentz矩阵P:")
print(P)

# 根据对称性列出方程，求解 cos theta 和 sin theta
print('对称性等式：')
expr1 = P[1,0] - P[0,1]
print(expr1) 
expr2 = P[2,0] - P[0,2]
print(expr2)
expr3 = P[1,2] - P[2,1]
print(expr3)

solution = sp.solve((expr1, expr3), (s, c))
print('sin:')
print(solution[s])
print(solution[s].subs(test_data_dict).evalf())
print('cos:')
print(solution[c])
print(solution[c].subs(test_data_dict).evalf())
